import torch
from d2l.torch import try_gpu, train_ch6
from torch import nn
from d2l import torch as d2l


def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, std=0.01)


def get_net():
    net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
    net.apply(init_weights)
    return net


net = get_net()
batch_size, lr, num_epochs = 256, 0.1, 100
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
# 使用GPU
train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
# 使用CPU
# train_ch6(net, train_iter, test_iter, num_epochs, lr, torch.device('cpu'))
# 保存训练得到的参数
torch.save(net.state_dict(), 'mlp.params')
# 加载模型的参数
net2 = get_net()
net2.load_state_dict(torch.load('mlp.params'))
# 查看第三层参数值
print(net2[3].state_dict())
# 测试集测试
d2l.predict_ch3(net2, test_iter)
d2l.plt.show()
